using System;
using System.Threading;
using System.Threading.Tasks;

namespace Sage.Http.Utils
{
    /// <summary>
    /// 令牌桶限流器，用于控制请求速率
    /// </summary>
    public class TokenBucket
    {
        private readonly int _capacity;
        private readonly double _refillRate;
        private readonly SemaphoreSlim _semaphore = new(1, 1);
        private double _tokens;
        private DateTime _lastRefillTime;

        /// <summary>
        /// 初始化令牌桶限流器
        /// </summary>
        /// <param name="capacity">桶容量（最大令牌数）</param>
        /// <param name="refillRate">每秒补充的令牌数</param>
        public TokenBucket(int capacity, double refillRate)
        {
            _capacity = capacity > 0 ? capacity : throw new ArgumentOutOfRangeException(nameof(capacity), "容量必须大于0");
            _refillRate = refillRate > 0 ? refillRate : throw new ArgumentOutOfRangeException(nameof(refillRate), "补充速率必须大于0");
            _tokens = capacity;
            _lastRefillTime = DateTime.UtcNow;
        }

        /// <summary>
        /// 尝试获取指定数量的令牌
        /// </summary>
        /// <param name="count">要获取的令牌数</param>
        /// <returns>是否成功获取令牌</returns>
        public bool TryTake(int count = 1)
        {
            if (count <= 0)
                throw new ArgumentOutOfRangeException(nameof(count), "令牌数必须大于0");

            _semaphore.Wait();
            try
            {
                RefillTokens();
                if (_tokens < count)
                {
                    return false;
                }

                _tokens -= count;
                return true;
            }
            finally
            {
                _semaphore.Release();
            }
        }

        /// <summary>
        /// 异步尝试获取指定数量的令牌
        /// </summary>
        /// <param name="count">要获取的令牌数</param>
        /// <param name="cancellationToken">取消令牌</param>
        /// <returns>是否成功获取令牌</returns>
        public async Task<bool> TryTakeAsync(int count = 1, CancellationToken cancellationToken = default)
        {
            if (count <= 0)
                throw new ArgumentOutOfRangeException(nameof(count), "令牌数必须大于0");

            await _semaphore.WaitAsync(cancellationToken);
            try
            {
                RefillTokens();
                if (_tokens < count)
                {
                    return false;
                }

                _tokens -= count;
                return true;
            }
            finally
            {
                _semaphore.Release();
            }
        }

        /// <summary>
        /// 异步等待并获取指定数量的令牌
        /// </summary>
        /// <param name="count">要获取的令牌数</param>
        /// <param name="timeout">超时时间</param>
        /// <param name="cancellationToken">取消令牌</param>
        /// <returns>是否成功获取令牌</returns>
        public async Task<bool> WaitAndTakeAsync(int count = 1, TimeSpan? timeout = null, CancellationToken cancellationToken = default)
        {
            if (count <= 0)
                throw new ArgumentOutOfRangeException(nameof(count), "令牌数必须大于0");

            var waitTime = timeout ?? TimeSpan.FromMilliseconds(-1); // -1表示无限等待
            var waitTask = WaitForTokensAsync(count, cancellationToken);

            if (waitTime.TotalMilliseconds < 0)
            {
                // 无限等待
                await waitTask;
                return true;
            }
            else
            {
                // 有超时限制
                if (await Task.WhenAny(waitTask, Task.Delay(waitTime, cancellationToken)) == waitTask)
                {
                    await waitTask; // 确保任何异常被抛出
                    return true;
                }
                else
                {
                    return false; // 超时
                }
            }
        }

        private async Task WaitForTokensAsync(int count, CancellationToken cancellationToken)
        {
            while (true)
            {
                await _semaphore.WaitAsync(cancellationToken);
                try
                {
                    RefillTokens();
                    if (_tokens >= count)
                    {
                        _tokens -= count;
                        return;
                    }
                }
                finally
                {
                    _semaphore.Release();
                }

                // 计算需要等待的时间
                double tokensNeeded = count - _tokens;
                double secondsToWait = tokensNeeded / _refillRate;
                int millisecondsToWait = (int)(secondsToWait * 1000);

                // 等待一段时间，但不要等待太久，以便能够响应取消
                await Task.Delay(Math.Min(millisecondsToWait, 100), cancellationToken);
            }
        }

        private void RefillTokens()
        {
            var now = DateTime.UtcNow;
            var elapsedSeconds = (now - _lastRefillTime).TotalSeconds;
            if (elapsedSeconds > 0)
            {
                var tokensToAdd = elapsedSeconds * _refillRate;
                _tokens = Math.Min(_capacity, _tokens + tokensToAdd);
                _lastRefillTime = now;
            }
        }
    }
}